package cn.denvie.agent.common;

import javassist.CannotCompileException;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.CtNewMethod;

import java.io.IOException;

/**
 * 字节码插桩工具类。
 *
 * @author DengZhaoyong
 * @version 1.0.0
 */
public class ByteAgentBuilder {

    private final String className;
    private final ClassLoader classLoader;
    private final CtClass ctClass;

    public ByteAgentBuilder(String className, ClassLoader classLoader, CtClass ctClass) {
        this.className = className;
        this.classLoader = classLoader;
        this.ctClass = ctClass;
    }

    /**
     * 更新方法，插入监听代码。
     */
    public void updateMethod(CtMethod method, MethodBuilder srcBuilder) throws CannotCompileException {
        CtMethod ctMethod = method;
        String methodName = method.getName();

        // 基于原方法复制生成代理方法
        CtMethod agentMethod = CtNewMethod.copy(ctMethod, methodName, ctClass, null);
        agentMethod.setName(methodName + "$Agent");
        ctClass.addMethod(agentMethod);

        // 原方法重置为代理执行
        ctMethod.setBody(srcBuilder.buildSrc(ctMethod));
    }

    /**
     * 生成新的class 字节码。
     */
    public byte[] toBytecote() throws IOException, CannotCompileException {
        return ctClass.toBytecode();
    }

}
